from core.text2image import text2image, text2image_once, text2imageN
from core.image2video import image2video
import streamlit as st
from PIL import Image
from typing import List

image_model_id = "../../hf/stable-diffusion-v1-5"
        
def main():
  # gen_image_ui()
  gen_video_ui()
  # test_gen_image()
  
def gen_image_ui():
  st.set_page_config(page_title="文生图测试", page_icon="📖=>🏞️")
  st.title("文生图测试")
  with st.form(key="my_form"):
    query = st.text_area(
      label="输入提示词..",
      help= "这里输入生成图片的提示词.",
      key="query",
      max_chars = 50,
    )
    submit_btn = st.form_submit_button(label="生成图片")
    
    if query and submit_btn:
      with st.spinner(text="正在生成图片..."):
        # prompt = "a photo of an astronaut riding a horse on mars"
        images = text2imageN(image_model_id, prompt = query, n=4)
        # images = [text2image(image_model_id, prompt = query) for _ in range(4)]
        # image = text2image(image_model_id, prompt = query)
        # image.save("astronaut_rides_horse.png")
      st.subheader("您生成的图片")
      # st.image(image= image, caption= query)
      # st.image(image= images, width=100, use_column_width=True, caption= [query]*4)
      grid_image_ui(2, images)
      
      

def grid_image_ui(n_cols: int, images:List):
  n_rows = 1 + len(images) // int(n_cols)
  rows = [st.container() for _ in range(n_rows)]
  cols_per_row = [r.columns(n_cols) for r in rows]
  cols = [column for row in cols_per_row for column in row]

  for image_index, image in enumerate(images):
      cols[image_index].image(image)

def gen_video_ui():
  st.set_page_config(page_title="文->图->视频测试", page_icon="📖=>🏞️")
  st.title("文->图->视频测试")
  with st.form(key="my_form"):
    query = st.text_area(
      label="输入提示词..",
      help= "这里输入生成图片的提示词.",
      key="query",
      max_chars = 50,
    )
    submit_btn = st.form_submit_button(label="生成视频")
    
    if query and submit_btn:
      with st.spinner(text="正在生成视频..."):
        # prompt = "a photo of an astronaut riding a horse on mars"
        image = text2image_once(image_model_id, prompt = query)
        print(f"image=>", image)
        # images = [text2image(image_model_id, prompt) for _ in range(4)]
        # image = Image.open("astronaut_rides_horse_1.png")
        video = image2video(query, image)
        # image.save("astronaut_rides_horse.png")
      st.subheader("您生成的video")
      st.video(data=video)


if __name__ == "__main__":
  main()
 